import numpy as np
import tensorflow as tf
import re
import ast
import traceback
import autopep8  # 新增：自动格式化工具
from tensorflow.keras.models import Sequential, load_model
from tensorflow.keras.layers import LSTM, Dense, Embedding
from tensorflow.keras.preprocessing.text import Tokenizer, tokenizer_from_json
from tensorflow.keras.preprocessing.sequence import pad_sequences
import json
from datetime import datetime


class CodeGeneratorAI:
    def __init__(self, max_sequence_length=200, vocab_size=5000):
        self.max_sequence_length = max_sequence_length
        self.vocab_size = vocab_size
        self.tokenizer = Tokenizer(num_words=vocab_size, filters='', lower=False)
        self.model = None
        self.error_patterns = {
            'SyntaxError': r'SyntaxError: (.*?) at line (\d+)',
            'IndentationError': r'IndentationError: (.*?) at line (\d+)',
            'NameError': r'NameError: name \'(\w+)\' is not defined',
            'TypeError': r'TypeError: (.*?) at line (\d+)',
            'IndexError': r'IndexError: (.*?) at line (\d+)'
        }

    def preprocess_code(self, code_samples):
        processed = []
        for code in code_samples:
            code = re.sub(r'#.*', '', code)
            code = code.replace('\t', '    ')
            code = code.strip('\n')
            try:
                ast.parse(code)
            except SyntaxError:
                continue
            processed.append(f'<START> {code} <END>')
        return processed

    def train(self, code_samples, epochs=150, batch_size=32):
        processed_code = self.preprocess_code(code_samples)
        if not processed_code:
            raise ValueError("有效训练数据为空")
        
        self.tokenizer.fit_on_texts(processed_code)
        total_words = len(self.tokenizer.word_index) + 1
        
        input_sequences = []
        for line in processed_code:
            token_list = self.tokenizer.texts_to_sequences([line])[0]
            for i in range(1, len(token_list)):
                n_gram_sequence = token_list[:i+1]
                input_sequences.append(n_gram_sequence)
        
        input_sequences = pad_sequences(
            input_sequences, 
            maxlen=self.max_sequence_length, 
            padding='pre',
            truncating='pre'
        )
        
        predictors, label = input_sequences[:,:-1], input_sequences[:,-1]
        label = tf.keras.utils.to_categorical(label, num_classes=total_words)
        
        self.model = Sequential()
        self.model.add(Embedding(total_words, 128, input_length=self.max_sequence_length-1, mask_zero=True))
        self.model.add(LSTM(200, return_sequences=True))
        self.model.add(LSTM(150))
        self.model.add(Dense(total_words, activation='softmax'))
        
        self.model.compile(loss='categorical_crossentropy', optimizer='adam', metrics=['accuracy'])
        self.model.fit(predictors, label, epochs=epochs, batch_size=batch_size, verbose=1)

    def generate_code(self, prompt, max_length=200, temperature=0.6):
        if self.model is None:
            raise Exception("模型尚未训练")
            
        generated = f'<START> {prompt}'
        for _ in range(max_length):
            token_list = self.tokenizer.texts_to_sequences([generated])[0]
            token_list = pad_sequences([token_list], maxlen=self.max_sequence_length-1, padding='pre')
            
            predictions = self.model.predict(token_list, verbose=0)[0]
            predictions = np.log(predictions + 1e-10) / temperature
            exp_preds = np.exp(predictions)
            predictions = exp_preds / np.sum(exp_preds)
            predicted_index = np.random.choice(len(predictions), p=predictions)
            
            output_word = ""
            for word, index in self.tokenizer.word_index.items():
                if index == predicted_index:
                    output_word = word
                    break
                    
            if output_word == '<END>':
                break
                
            generated += " " + output_word
            
        return self._clean_tokens(generated)

    def analyze_error(self, code, error_message):
        error_info = {'type': 'Unknown', 'line': -1, 'details': ''}
        
        for error_type, pattern in self.error_patterns.items():
            match = re.search(pattern, error_message)
            if match:
                error_info['type'] = error_type
                if error_type in ['SyntaxError', 'IndentationError', 'TypeError', 'IndexError']:
                    error_info['line'] = int(match.group(2))
                    error_info['details'] = match.group(1)
                elif error_type == 'NameError':
                    error_info['undefined_name'] = match.group(1)
                break
        
        lines = code.split('\n')
        if 0 <= error_info['line'] - 1 < len(lines):
            error_info['error_line'] = lines[error_info['line'] - 1]
            start = max(0, error_info['line'] - 4)
            end = min(len(lines), error_info['line'] + 3)
            error_info['context'] = lines[start:end]
        
        return error_info

    def fix_code(self, code, error_info, max_length=150, temperature=0.5):
        prompt = self._build_fix_prompt(code, error_info)
        fixed_code = self.generate_code(prompt, max_length=max_length, temperature=temperature)
        
        try:
            ast.parse(fixed_code)
            return fixed_code
        except SyntaxError:
            # 尝试使用autopep8格式化（如果只是格式问题）
            try:
                return self._format_code(fixed_code)
            except:
                return code

    def refine_code(self, code, instruction, context=None, max_length=200, temperature=0.6):
        prompt = self._build_refine_prompt(code, instruction, context)
        refined_code = self.generate_code(prompt, max_length=max_length, temperature=temperature)
        
        # 使用autopep8格式化优化后的代码
        return self._format_code(refined_code)

    def _format_code(self, code):
        """使用autopep8格式化代码"""
        try:
            # 严格遵守PEP8规范，修正缩进、空格等问题
            return autopep8.fix_code(
                code, 
                options={
                    'aggressive': 2,  # 更激进的格式化
                    'ignore': ['E501']  # 忽略行长度限制
                }
            )
        except Exception as e:
            print(f"格式化失败: {e}")
            return code  # 格式化失败时返回原始代码

    def _build_refine_prompt(self, code, instruction, context=None):
        prompt = "# 当前代码:\n" + code + "\n\n"
        
        if context:
            prompt += "# 历史对话上下文:\n"
            for message in context:
                role = message['role']
                content = message['content']
                prompt += f"# [{role}]: {content}\n"
            prompt += "\n"
        
        prompt += f"# 新指令: {instruction}\n"
        prompt += "# 请根据以上信息优化代码:"
        
        return prompt

    def _build_fix_prompt(self, code, error_info):
        prompt = f"""
# 原始代码（有错误）
{code}

# 错误信息
{error_info['type']}: {error_info.get('details', '')}
{'错误行: ' + str(error_info['line']) if error_info['line'] != -1 else ''}

# 修复指令
请修复上述代码中的{error_info['type']}错误。
{"特别注意第" + str(error_info['line']) + "行" if error_info['line'] != -1 else ""}
{"确保定义了" + error_info['undefined_name'] if 'undefined_name' in error_info else ""}
"""
        return prompt.strip()

    def _clean_tokens(self, text):
        text = text.replace('<START> ', '')
        text = text.replace(' <END>', '')
        text = text.replace('<START>', '')
        text = text.replace('<END>', '')
        text = re.sub(r'\n+', '\n', text)
        return text.strip()

    def save_model(self, model_path='code_model.keras', tokenizer_path='tokenizer_config.json'):
        self.model.save(model_path)
        with open(tokenizer_path, 'w') as f:
            f.write(self.tokenizer.to_json())

    def load_model(self, model_path='code_model.keras', tokenizer_path='tokenizer_config.json'):
        self.model = load_model(model_path)
        with open(tokenizer_path, 'r') as f:
            self.tokenizer = tokenizer_from_json(f.read())


# 增强版交互式代码编辑器
class InteractiveCodeEditor:
    def __init__(self, code_generator):
        self.code_generator = code_generator
        self.current_code = ""
        self.context = []
        self.history = []

    def generate_from_prompt(self, prompt, temperature=0.6):
        print(f"正在根据提示生成代码...")
        self.current_code = self.code_generator.generate_code(prompt, temperature=temperature)
        self.current_code = self.code_generator._format_code(self.current_code)  # 格式化生成的代码
        
        self.context.append({'role': 'user', 'content': prompt})
        self.context.append({'role': 'assistant', 'content': self.current_code})
        
        self.history.append(("生成", prompt, self.current_code))
        return self.current_code

    def refine_code(self, instruction, temperature=0.6):
        if not self.current_code:
            print("没有当前代码，请先生成代码")
            return
        
        print(f"正在根据指令优化代码...")
        self.current_code = self.code_generator.refine_code(
            self.current_code, 
            instruction, 
            context=self.context,
            temperature=temperature
        )
        
        self.context.append({'role': 'user', 'content': instruction})
        self.context.append({'role': 'assistant', 'content': self.current_code})
        
        self.history.append(("优化", instruction, self.current_code))
        return self.current_code

    def fix_error(self, error_message, temperature=0.5):
        if not self.current_code:
            print("没有当前代码，请先生成代码")
            return
        
        print(f"正在分析错误并修复代码...")
        error_info = self.code_generator.analyze_error(self.current_code, error_message)
        
        print(f"错误类型: {error_info['type']}")
        if error_info['line'] != -1:
            print(f"错误位置: 第{error_info['line']}行")
        if 'undefined_name' in error_info:
            print(f"未定义的变量: {error_info['undefined_name']}")
        if 'error_line' in error_info:
            print(f"错误行: {error_info['error_line']}")
        
        # 尝试AI修复
        self.current_code = self.code_generator.fix_code(self.current_code, error_info, temperature=temperature)
        
        # 再次使用autopep8格式化
        self.current_code = self.code_generator._format_code(self.current_code)
        
        self.context.append({'role': 'user', 'content': f"错误: {error_message}"})
        self.context.append({'role': 'assistant', 'content': self.current_code})
        
        self.history.append(("修复", error_message, self.current_code))
        return self.current_code

    def auto_format(self):
        """使用autopep8自动格式化当前代码"""
        if not self.current_code:
            print("没有当前代码，请先生成代码")
            return
        
        original_code = self.current_code
        self.current_code = self.code_generator._format_code(self.current_code)
        
        if original_code != self.current_code:
            print("代码已自动格式化，符合PEP8规范")
            self.context.append({'role': 'user', 'content': "自动格式化代码"})
            self.context.append({'role': 'assistant', 'content': self.current_code})
            self.history.append(("格式化", "自动PEP8格式化", self.current_code))
        else:
            print("代码已经符合PEP8规范")
        
        return self.current_code

    def check_pep8(self):
        """检查代码的PEP8合规性（返回问题列表）"""
        if not self.current_code:
            print("没有当前代码，请先生成代码")
            return []
        
        try:
            # 使用autopep8的检查功能
            options = autopep8.parse_args(['--diff', '--aggressive', '--ignore=E501'])
            diff = autopep8.fix_string(self.current_code, options)
            
            if diff == self.current_code:
                print("代码完全符合PEP8规范")
                return []
            
            # 解析diff输出，提取问题
            issues = []
            lines = diff.split('\n')
            for line in lines:
                if line.startswith('@@'):
                    # 提取行号信息
                    parts = line.split()
                    if len(parts) >= 3:
                        line_info = parts[2]
                        if ',' in line_info:
                            line_num = int(line_info.split(',')[0].lstrip('-'))
                            issues.append(f"行 {line_num}: 格式问题")
            
            if issues:
                print(f"发现 {len(issues)} 个PEP8格式问题")
            return issues
            
        except Exception as e:
            print(f"检查失败: {e}")
            return []

    def edit_code(self, new_code):
        self.current_code = new_code
        self.context.append({'role': 'user', 'content': "手动编辑代码"})
        self.context.append({'role': 'assistant', 'content': self.current_code})
        self.history.append(("手动编辑", "用户修改", self.current_code))
        return self.current_code

    def show_context(self):
        print("\n=== 对话上下文 ===")
        for i, message in enumerate(self.context):
            role = message['role']
            content = message['content']
            preview = content[:50] + ('...' if len(content) > 50 else '')
            print(f"{i+1}. [{role}]: {preview}")
        print("==================\n")

    def save_code(self, filename=None):
        if not filename:
            filename = f"generated_code_{datetime.now().strftime('%Y%m%d_%H%M%S')}.py"
            
        with open(filename, 'w') as f:
            f.write(self.current_code)
            
        print(f"代码已保存到: {filename}")
        return filename

    def show_history(self):
        print("\n=== 操作历史 ===")
        for i, (action, instruction, _) in enumerate(self.history):
            preview = instruction[:50] + ('...' if len(instruction) > 50 else '')
            print(f"{i+1}. [{action}] {preview}")
        print("================\n")

    def reset_context(self):
        self.context = []
        print("对话上下文已重置")


# 示例用法
if __name__ == "__main__":
    # 训练数据（增加格式化相关示例）
    training_data = [
        "def add(a, b):\n    return a + b",
        "# 未格式化的代码\nx=5\ny=10\nprint(x+y)\n\n# 格式化后的代码\nx = 5\ny = 10\nprint(x + y)",
        "# 初始代码\nresult = add(3,4)\n\n# 指令: 添加空格，符合PEP8\n# 优化后的代码\nresult = add(3, 4)",
        "# 当前代码\nif x>5:\n    print('大')\n\n# 指令: 修复缩进和空格\n# 优化后的代码\nif x > 5:\n    print('大')"
    ]
    
    # 初始化模型
    ai = CodeGeneratorAI()
    ai.train(training_data, epochs=50)
    
    # 初始化编辑器
    editor = InteractiveCodeEditor(ai)
    
    # 示例：生成不规范代码
    prompt = "def calculate_average(numbers):\n    if not numbers: return 0\n    return sum(numbers)/len(numbers)"
    editor.generate_from_prompt(prompt)
    
    # 示例：自动格式化
    editor.auto_format()
    
    # 示例：检查PEP8合规性
    issues = editor.check_pep8()
    print(f"PEP8问题: {issues}")
    
    # 示例：显示最终代码
    print("\n最终代码：")
    print(editor.current_code)